import os

import nibabel as nib
import numpy as np
from scipy import ndimage
from scipy.spatial import distance
from skimage import measure


#####################################################
# Pre-define parameters
#####################################################

def define_parameter():
    params = {}

    # ----------------------------------------------------
    # Parameters for intensity (fixed)
    # ----------------------------------------------------

    params['lungMinValue'] = -1024
    params['lungMaxValue'] = -400
    # params['lungThreshold'] = -900
    params['lungThreshold'] = -359

    # ----------------------------------------------------
    # Parameters for lung segmentation (fixed)
    # ----------------------------------------------------

    params['xRangeRatio1'] = 0.4
    params['xRangeRatio2'] = 0.75
    params['zRangeRatio1'] = 0.5
    params['zRangeRatio2'] = 0.75

    # ----------------------------------------------------
    # Parameters for airway segmentation
    # NEED TO ADAPT for image resolution and orientation
    # [current values work for demo image with resolution = 1mm^3]
    # ----------------------------------------------------
    params['airwayRadiusMask'] = 15  # increase the value if you have high resolution image
    params['airwayRadiusX'] = 8  # ditto
    params['airwayRadiusZ'] = 15  # ditto
    params['super2infer'] = 0  # value = 1 if slice no. increases from superior to inferior, else value = 0

    params['currentSampleRootPath'] = ""
    params['currentSampleName'] = ""

    return params


#####################################################
# Generate binary structure to mimic trachea
#####################################################

def generate_structure_trachea(Radius, RadiusX, RadiusZ):
    struct_trachea = np.zeros([2 * Radius + 1, 2 * Radius + 1, RadiusZ])
    for i in range(0, 2 * Radius + 1):
        for j in range(0, 2 * Radius + 1):
            if distance.euclidean([Radius + 1, Radius + 1], [i, j]) < RadiusX:
                struct_trachea[i, j, :] = 1
            else:
                struct_trachea[i, j, :] = 0

    return struct_trachea


#####################################################
# Generate bounding box
#####################################################

def bbox2_3D(img, label, margin, limit):
    imgtmp = np.zeros(img.shape)
    imgtmp[img == label] = 1

    x = np.any(imgtmp, axis=(1, 2))
    y = np.any(imgtmp, axis=(0, 2))
    z = np.any(imgtmp, axis=(0, 1))

    xmin, xmax = np.where(x)[0][[0, -1]]
    ymin, ymax = np.where(y)[0][[0, -1]]
    zmin, zmax = np.where(z)[0][[0, -1]]

    xmin = xmin - margin - 1
    xmin = max(0, xmin)
    ymin = ymin - margin - 1
    ymin = max(0, ymin)
    zmin = zmin - margin - 1
    zmin = max(0, zmin)
    xmax = xmax + margin + 1
    xmax = min(xmax, limit[0])
    ymax = ymax + margin + 1
    ymax = min(ymax, limit[1])
    zmax = zmax + margin + 1
    zmax = min(zmax, limit[2])

    return xmin, xmax, ymin, ymax, zmin, zmax


#####################################################
# Generate inital point in trachea
#####################################################

def generate_initLoc(params, I, Mlung, Radius, RadiusZ, struct_trachea):
    mind = np.argwhere(Mlung == 1)
    initLoc = [0, 0, 0]
    minDiff = float('inf')

    if params['super2infer']:
        slice_no = np.min(mind[:, 2])
        Itmp = I[:, :, slice_no:slice_no + RadiusZ]
    else:
        slice_no = np.max(mind[:, 2])
        Itmp = I[:, :, slice_no - RadiusZ:slice_no]

    Mtmp = np.ones(Itmp.shape);
    Mtmp[Itmp < params['lungMinValue']] = 0
    Mtmp[Itmp > params['lungMaxValue']] = 0
    Itmp = Mtmp;
    Mtmp = np.sum(Mtmp, axis=2)

    for i in range(Radius, Itmp.shape[0] - Radius):
        for j in range(Radius, Itmp.shape[1] - Radius):
            if Mtmp[i, j] > 0:
                struct_Itmp = Itmp[i - Radius:i + Radius + 1, j - Radius:j + Radius + 1, :]
                currVal = struct_Itmp - struct_trachea
                currVal = np.sum(np.square(currVal))

                if currVal < minDiff:
                    initLoc = [i, j, slice_no]
                    minDiff = currVal

    print('initial location = ' + str(initLoc))

    return slice_no, initLoc


#####################################################
# Closed space dialation to segment airway
#####################################################

def close_space_dilation(params, I, Mlung, Radius, RadiusX, RadiusZ, struct_s, slice_no, initLoc):
    """

    :param params:
    :param I: I is the original image
    :param Mlung: Mlung is the result of lung segmentation
    :param Radius:
    :param RadiusX:
    :param RadiusZ:
    :param struct_s:
    :param slice_no:
    :param initLoc:
    :return: the airway
    """
    iterNoPerSlice = RadiusX
    maxFactor = RadiusX / 2
    maxChange = RadiusX * RadiusX * RadiusX * 50
    totalChange = 1
    tempCheck = 0
    [m, n, p] = I.shape
    Mtmp = np.zeros([m, n, p])
    struct_m = ndimage.iterate_structure(struct_s, 2)

    if params['super2infer']:
        Mtmp[initLoc[0] - Radius:initLoc[0] + Radius + 1,
        initLoc[1] - Radius:initLoc[1] + Radius + 1,
        0:slice_no + RadiusZ] = 1
    else:
        Mtmp[initLoc[0] - Radius:initLoc[0] + Radius + 1,
        initLoc[1] - Radius:initLoc[1] + Radius + 1,
        slice_no - RadiusZ:p - 1] = 1
    Mtmp = np.multiply(Mtmp, Mlung)
    # save_nii_file(Mtmp, "Mtmp")
    Minit = ndimage.binary_closing(Mtmp, structure=struct_s, iterations=1)
    Minit = np.int8(Minit)
    Minit[Minit > 0] = 2

    while totalChange > 0:

        maxSegmentChange = 0
        tempCheck = tempCheck + 1
        L = measure.label(np.floor(Minit / 2))
        Minit[Minit > 1] = 1

        for label in np.unique(L[:]):

            if label != 0 and np.sum(L[:] == label) > 10:

                # Process each component in local FOV

                xmin, xmax, ymin, ymax, zmin, zmax = bbox2_3D(L, label, iterNoPerSlice, [m, n, p])
                Mtmp = Minit[xmin:xmax, ymin:ymax, zmin:zmax]
                Itmp = I[xmin:xmax, ymin:ymax, zmin:zmax]
                Ltmp = L[xmin:xmax, ymin:ymax, zmin:zmax]
                Ltmp[Ltmp != label] = 0
                Ltmp[Ltmp > 0] = 1;

                for iterCount in range(0, iterNoPerSlice):
                    Ltmp = ndimage.binary_dilation(Ltmp, structure=struct_s, iterations=1)
                    Ltmp = np.int8(Ltmp)
                    Ltmp[Itmp > params['lungThreshold']] = 0

                Ltmp = ndimage.binary_closing(Ltmp, structure=struct_s, iterations=1)
                Ltmp = np.int8(Ltmp)
                Ltmp[Mtmp > 0] = 0
                Ltmp[Ltmp > 0] = 2
                Ltmp = Ltmp + Mtmp

                segmentChange = np.sum(Ltmp[:] > 1)
                if segmentChange < maxChange or tempCheck < 10:
                    Minit[xmin:xmax, ymin:ymax, zmin:zmax] = Ltmp
                    if segmentChange > maxSegmentChange:
                        maxSegmentChange = segmentChange

        if tempCheck < 10:
            maxChange = max(maxFactor * maxSegmentChange, maxChange)
        else:
            maxChange = min(maxFactor * maxSegmentChange, maxChange)

        totalChange = np.sum(Minit[:] > 1)

        print('iter = ' + str(tempCheck) + ' airway sum = ' + str(np.sum(Minit[:] > 0)) \
              + ' airway change = ' + str(totalChange))

    Minit[Minit > 0] = 1
    Minit = ndimage.binary_opening(Minit, structure=struct_s, iterations=1)
    Minit = ndimage.binary_dilation(Minit, structure=struct_m, iterations=1)
    Maw = np.int8(Minit)

    return Maw


def save_nii_file(output_image, output_filename,
                  sample_file_path=r"D:\dataset\vessel\Parse_2022_train_data\PA000005\image\PA000005.nii.gz"):
    # output_image = output_image.transpose((2, 1, 0))
    if isinstance(output_image[0][0][0].item(), bool):
        output_image = np.where(output_image == True, 1, 0).astype(dtype="int8")
    func = nib.load(sample_file_path)
    ni_img = nib.Nifti1Image(output_image, func.affine)
    output_path = os.path.join(os.path.dirname(sample_file_path), output_filename + ".nii.gz")
    nib.save(ni_img, output_path)
    print(output_filename + " has been written.")
    print("The output path is " + output_path)
